#pragma region Copyright
/// <copyright>
/// Copyright (c) 2007 Ilia Shramko.
/// Released under the terms of the GNU General Public License.
/// </copyright>
#pragma endregion

#include "stdafx.h"

const wchar_t* SavedInprocServer32PathKey = L"DefaultSaved";
const wchar_t* CLSIDKeyName = L"CLSID\\";
const wchar_t* InprocServer32KeyName = L"\\InprocServer32";

Module::Module() : _dxLibrary(NULL), _dllGetClassObject(NULL)
{
}

Module::~Module()
{
	FreeDXLibrary();
}

HRESULT Module::DllRegisterServer(BOOL bRegTypeLib)
{
	HRESULT hr = CAtlDllModuleT<Module>::DllRegisterServer(bRegTypeLib);
	if(FAILED(hr)) return hr;

	wchar_t moduleFilePath[MAX_PATH];
	DWORD len = ::GetModuleFileName(_AtlBaseModule.GetModuleInstance(), moduleFilePath, MAX_PATH);
	if(len == 0) return AtlHresultFromLastError();

	HRESULT hr1 = RegisterClass(CLSID_FilterGraph, moduleFilePath);
	HRESULT hr2 = RegisterClass(CLSID_FilterGraphNoThread, moduleFilePath);
	if(FAILED(hr1)) return hr1;
	return hr2;
}

HRESULT Module::DllUnregisterServer(BOOL bUnRegTypeLib)
{
	HRESULT hr = CAtlDllModuleT<Module>::DllUnregisterServer(bUnRegTypeLib);
	if(FAILED(hr)) return hr;

	HRESULT hr1 = UnregisterClass(CLSID_FilterGraph);
	HRESULT hr2 = UnregisterClass(CLSID_FilterGraphNoThread);
	if(FAILED(hr1)) return hr1;
	return hr2;
}

HRESULT Module::RegisterClass(REFCLSID rclsid, const wchar_t* modulePath)
{
	if(modulePath == NULL) return E_INVALIDARG;

	wchar_t clsidKeyPath[100] = {0};
	wcscpy_s(clsidKeyPath, _countof(clsidKeyPath), CLSIDKeyName);
	StringFromGUID2(rclsid, clsidKeyPath + wcslen(clsidKeyPath), 50);
	wcscat_s(clsidKeyPath, _countof(clsidKeyPath), InprocServer32KeyName);
	
	HRESULT hr = Security::SetRequiredPermissionsToRegistryKey(HKEY_CLASSES_ROOT, clsidKeyPath);
	if(FAILED(hr)) return hr;

	ATL::CRegKey clsidKey;
	LONG result = clsidKey.Open(HKEY_CLASSES_ROOT, clsidKeyPath);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	// Get source InprocServer32 path of rclsid class (path to quartz.dll)
	wchar_t sourcePath[MAX_PATH + 1];
	ULONG sourcePathLength = MAX_PATH;
	result = clsidKey.QueryStringValue(NULL, sourcePath, &sourcePathLength);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	// Check if InprocServer32 key was saved earlier.
	wchar_t savedSourcePath[MAX_PATH + 1] = {0};
	ULONG savedSourcePathLength = MAX_PATH;
	result = clsidKey.QueryStringValue(SavedInprocServer32PathKey, savedSourcePath, &savedSourcePathLength);
	if(result != ERROR_SUCCESS || savedSourcePathLength == 0)
	{
		// Save source InprocServer32 path of rclsid class if it was not saved earlier.
		result = clsidKey.SetStringValue(SavedInprocServer32PathKey, sourcePath);
		if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);
	}

	// Change InprocServer32 path of rclsid class to point to modulePath
	result = clsidKey.SetStringValue(NULL, modulePath);
	return AtlHresultFromWin32(result);
}

HRESULT Module::UnregisterClass(REFCLSID rclsid)
{
	wchar_t clsidKeyPath[100] = {0};
	wcscpy_s(clsidKeyPath, _countof(clsidKeyPath), CLSIDKeyName);
	StringFromGUID2(rclsid, clsidKeyPath + wcslen(clsidKeyPath), 50);
	wcscat_s(clsidKeyPath, _countof(clsidKeyPath), InprocServer32KeyName);
	
	HRESULT hr = Security::SetRequiredPermissionsToRegistryKey(HKEY_CLASSES_ROOT, clsidKeyPath);
	if(FAILED(hr)) return hr;

	ATL::CRegKey clsidKey;
	LONG result = clsidKey.Open(HKEY_CLASSES_ROOT, clsidKeyPath);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	// Read saved InprocServer32 value.
	wchar_t savedSourcePath[MAX_PATH + 1] = {0};
	ULONG savedSourcePathLength = MAX_PATH;
	result = clsidKey.QueryStringValue(SavedInprocServer32PathKey, savedSourcePath, &savedSourcePathLength);
	if(result != ERROR_SUCCESS) return S_OK;

	if(savedSourcePathLength == 0) return E_FAIL;

	// Restore saved InprocServer32 key value.
	result = clsidKey.SetStringValue(NULL, savedSourcePath);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	// Delete saved InprocServer32 value.
	result = clsidKey.DeleteValue(SavedInprocServer32PathKey);
	return AtlHresultFromWin32(result);
}

HRESULT Module::DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
{
	HRESULT hr = CAtlDllModuleT<Module>::DllGetClassObject(rclsid, riid, ppv);
	if(SUCCEEDED(hr)) return hr;
	if(hr != CLASS_E_CLASSNOTAVAILABLE) return hr;

	if(rclsid == CLSID_FilterGraph || rclsid == CLSID_FilterGraphNoThread)
	{
		if(ppv == NULL) return E_POINTER;
		*ppv = NULL;

		hr = LoadDXLibrary();

		ATLTRACE_TS(atlTraceGeneral, TraceLevel::Info, "Module::DllGetClassObject(%s) : 0x%08X\r\n",
			(rclsid == CLSID_FilterGraph ? "CLSID_FilterGraph" : "CLSID_FilterGraphNoThread"), hr);

		if(FAILED(hr)) return CLASS_E_CLASSNOTAVAILABLE;

		CComPtr<IClassFactory> classFactory;
		hr = (*_dllGetClassObject)(rclsid, __uuidof(IClassFactory), (void**)&classFactory);
		if(FAILED(hr)) return CLASS_E_CLASSNOTAVAILABLE;


		// Check application is enabled.
		bool appEnabled = false;
		HRESULT hr = Configuration::Current().IsCurrentAppEnabled();
		if(FAILED(hr) || hr == S_FALSE)
		{
			ATLTRACE_TS(atlTraceGeneral, TraceLevel::Info,
				"Module::DllGetClassObject() : application %S disabled - 0x%08X.\r\n",
				Configuration::Current().GetCurrentAppName().GetString(), hr);
		}
		else
		{
			appEnabled = true;
			ATLTRACE_TS(atlTraceGeneral, TraceLevel::Info,
				"Module::DllGetClassObject() : application %S enabled.\r\n",
				Configuration::Current().GetCurrentAppName().GetString());
		}


		// If current application should be ignored, return source class factory.
		if(!appEnabled)
		{
			*ppv = classFactory.Detach();
			return S_OK;
		}


		// Create wrapper
		//FilterGraphClassFactory* filterGraphClassFactory = NULL;
		CComObject<FilterGraphClassFactory>* filterGraphClassFactory;
		hr = CComObject<FilterGraphClassFactory>::CreateInstance(&filterGraphClassFactory);
		if(FAILED(hr)) return hr;
		
		filterGraphClassFactory->AddRef();

		hr = filterGraphClassFactory->Initialize(classFactory);
		if(SUCCEEDED(hr))
			hr = filterGraphClassFactory->QueryInterface(riid, ppv);

		filterGraphClassFactory->Release();
	}

	return hr;
}

HRESULT Module::LoadDXLibrary()
{
	if(_dxLibrary != NULL)
		return S_OK;

	CString dxLibraryPath;
	HRESULT hr = GetDXLibraryPath(dxLibraryPath);
	if(FAILED(hr)) return hr;

	_dxLibrary = ::LoadLibrary(dxLibraryPath);
	if(_dxLibrary == NULL)
		return AtlHresultFromLastError();

	_dllGetClassObject = (DllGetClassObjectProc*)::GetProcAddress(_dxLibrary, "DllGetClassObject");
	if(_dllGetClassObject == NULL)
		hr = AtlHresultFromLastError();

	if(SUCCEEDED(hr)) return hr;

	FreeDXLibrary();
	return hr;
}

HRESULT Module::FreeDXLibrary()
{
	if(_dxLibrary == NULL)
		return S_OK;

	if(::FreeLibrary(_dxLibrary))
	{
		_dllGetClassObject = NULL;
		_dxLibrary = NULL;
		return S_OK;
	}

	return AtlHresultFromLastError();
}

HRESULT Module::GetDXLibraryPath(CString& dxLibraryPath)
{
	wchar_t clsidKeyPath[100] = {0};
	wcscpy_s(clsidKeyPath, _countof(clsidKeyPath), CLSIDKeyName);
	StringFromGUID2(CLSID_FilterGraph, clsidKeyPath + wcslen(clsidKeyPath), 50);
	wcscat_s(clsidKeyPath, _countof(clsidKeyPath), InprocServer32KeyName);
	
	ATL::CRegKey clsidKey;
	LONG result = clsidKey.Open(HKEY_CLASSES_ROOT, clsidKeyPath);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	// Read saved InprocServer32 value.
	wchar_t savedSourcePath[MAX_PATH + 1] = {0};
	ULONG savedSourcePathLength = MAX_PATH;
	result = clsidKey.QueryStringValue(SavedInprocServer32PathKey, savedSourcePath, &savedSourcePathLength);
	if(result != ERROR_SUCCESS) return AtlHresultFromWin32(result);

	if(savedSourcePathLength == 0) return E_FAIL;

	dxLibraryPath = savedSourcePath;
	return S_OK;
}



